import os
import torch
import numpy as np
import cv2
from PIL import Image
from torchvision import transforms
import networks.resnet_encoder as resnet_encoder
import networks.depth_decoder as depth_decoder
from layers import disp_to_depth

class DepthEstimator:
    def __init__(self, model_name,dirpath="models", use_cuda=True):
        self.device = torch.device("cuda" if torch.cuda.is_available() and use_cuda else "cpu")
        self.model_path = os.path.join(dirpath, model_name)
        self.encoder_path = os.path.join(self.model_path, "encoder.pth")
        self.depth_decoder_path = os.path.join(self.model_path, "depth.pth")
        self.load_model()

    def load_model(self):
        # Load the encoder and decoder models
        self.encoder = resnet_encoder.ResnetEncoder(18, None)
        loaded_dict_enc = torch.load(self.encoder_path, map_location=self.device,weights_only=True)
        self.feed_height = loaded_dict_enc['height']
        self.feed_width = loaded_dict_enc['width']
        filtered_dict_enc = {k: v for k, v in loaded_dict_enc.items() if k in self.encoder.state_dict()}
        self.encoder.load_state_dict(filtered_dict_enc)
        self.encoder.to(self.device)
        self.encoder.eval()
        self.depth_decoder = depth_decoder.DepthDecoder(
            num_ch_enc=self.encoder.num_ch_enc, scales=range(4))
        loaded_dict = torch.load(self.depth_decoder_path, map_location=self.device,weights_only=True)
        self.depth_decoder.load_state_dict(loaded_dict)
        self.depth_decoder.to(self.device)
        self.depth_decoder.eval()
    def get_depth(self, img):
        """
        Predict and return an OpenCV Mat format depth map data
        """
        input_image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        original_height, original_width = input_image.size
        input_image = input_image.resize((self.feed_width, self.feed_height), Image.LANCZOS)
        input_image = transforms.ToTensor()(input_image).unsqueeze(0).to(self.device)

        with torch.no_grad():
            features = self.encoder(input_image)
            outputs = self.depth_decoder(features)
            disp = outputs[("disp", 0)]
            disp_resized = torch.nn.functional.interpolate(
                disp, (original_height, original_width), mode="bilinear", align_corners=False)

            # Convert disparity to depth
            scaled_disp, depth = disp_to_depth(disp_resized, 0.1, 100.0)

        # Convert to OpenCV Mat format
        depth_map = depth.squeeze().cpu().numpy()
        depth_map = cv2.normalize(depth_map, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        h,w,_=img.shape
        depth_map = cv2.resize(depth_map, (w,h))
        return depth_map

if __name__ == "__main__":
    imgpath="test.png"
    img = cv2.imread(imgpath)
    
    depth_estimator = DepthEstimator("mono+stereo_640x192")
    depth_map = depth_estimator.get_depth(img)
    
    cv2.imshow("origin img", img)
    cv2.imshow("Depth Map", depth_map)
    cv2.waitKey(0)